{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "This example requires the following dependencies to be installed:\n",
    "pip install lightly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install lightly"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "Note: The model and training settings do not follow the reference settings\n",
    "from the paper. The settings are chosen such that the example can easily be\n",
    "run on a small dataset with a single GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "from functools import partial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytorch_lightning as pl\n",
    "import torch\n",
    "import torchvision\n",
    "from timm.models.vision_transformer import vit_small_patch16_224\n",
    "from torch import Tensor\n",
    "from torch.nn import Module\n",
    "from torch.optim import AdamW"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lightly.loss import DINOLoss, IBOTPatchLoss, KoLeoLoss\n",
    "from lightly.models.modules import DINOv2ProjectionHead, MaskedVisionTransformerTIMM\n",
    "from lightly.models.utils import (\n",
    "    random_block_mask,\n",
    "    update_drop_path_rate,\n",
    "    update_momentum,\n",
    ")\n",
    "from lightly.transforms.dino_transform import DINOTransform\n",
    "from lightly.utils.scheduler import cosine_schedule, linear_warmup_schedule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def freeze_eval_module(module: Module) -> None:\n",
    "    \"\"\"Freeze the parameters of a module.\"\"\"\n",
    "    for param in module.parameters():\n",
    "        param.requires_grad = False\n",
    "    module.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DINOv2Head(Module):\n",
    "    def __init__(\n",
    "        self, dino_head: DINOv2ProjectionHead, ibot_head: DINOv2ProjectionHead\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "        self.dino_head = dino_head\n",
    "        self.ibot_head = ibot_head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DINOv2(pl.LightningModule):\n",
    "    def __init__(\n",
    "        self,\n",
    "        ibot_separate_head: bool = False,\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "        self.save_hyperparameters()\n",
    "\n",
    "        # Backbones\n",
    "        vit_teacher = vit_small_patch16_224(\n",
    "            pos_embed=\"learn\",\n",
    "            dynamic_img_size=True,\n",
    "            init_values=1e-5,\n",
    "        )\n",
    "        self.teacher_backbone = MaskedVisionTransformerTIMM(\n",
    "            vit=vit_teacher,\n",
    "            antialias=False,\n",
    "            pos_embed_initialization=\"skip\",\n",
    "        )\n",
    "        self.student_backbone = copy.deepcopy(self.teacher_backbone)\n",
    "        update_drop_path_rate(\n",
    "            self.student_backbone.vit,\n",
    "            drop_path_rate=0.1,  # we recommend using smaller rates like 0.1 for vit-s-14\n",
    "            mode=\"uniform\",\n",
    "        )\n",
    "\n",
    "        freeze_eval_module(self.teacher_backbone)\n",
    "\n",
    "        # Heads\n",
    "        dino_head = partial(\n",
    "            DINOv2ProjectionHead,\n",
    "            input_dim=384,\n",
    "        )\n",
    "\n",
    "        teacher_dino_head = dino_head()\n",
    "        student_dino_head = dino_head()\n",
    "\n",
    "        ibot_head = partial(\n",
    "            DINOv2ProjectionHead,\n",
    "            input_dim=384,\n",
    "        )\n",
    "\n",
    "        if ibot_separate_head:\n",
    "            teacher_ibot_head = ibot_head()\n",
    "            student_ibot_head = ibot_head()\n",
    "        else:\n",
    "            teacher_ibot_head = teacher_dino_head\n",
    "            student_ibot_head = student_dino_head\n",
    "\n",
    "        self.teacher_head = DINOv2Head(\n",
    "            dino_head=teacher_dino_head,\n",
    "            ibot_head=teacher_ibot_head,\n",
    "        )\n",
    "        self.student_head = DINOv2Head(\n",
    "            dino_head=student_dino_head,\n",
    "            ibot_head=student_ibot_head,\n",
    "        )\n",
    "\n",
    "        freeze_eval_module(self.teacher_head)\n",
    "\n",
    "        # Losses\n",
    "        self.dino_criterion = DINOLoss()\n",
    "        self.ibot_criterion = IBOTPatchLoss()\n",
    "        self.koleo_criterion = KoLeoLoss()\n",
    "\n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        pass\n",
    "\n",
    "    def forward_teacher(self, x: Tensor) -> tuple[Tensor, Tensor]:\n",
    "        features = self.teacher_backbone.encode(x)\n",
    "        cls_tokens = features[:, 0]\n",
    "        return cls_tokens, features\n",
    "\n",
    "    def forward_student(\n",
    "        self, x: Tensor, mask: Tensor | None\n",
    "    ) -> tuple[Tensor, Tensor | None]:\n",
    "        features = self.student_backbone.encode(x, mask=mask)\n",
    "        cls_tokens = features[:, 0]\n",
    "        masked_features = None if mask is None else features[mask]\n",
    "        return cls_tokens, masked_features\n",
    "\n",
    "    def training_step(\n",
    "        self, batch: tuple[list[Tensor], Tensor, list[str]], batch_idx: int\n",
    "    ) -> Tensor:\n",
    "        views, _ = batch[0], batch[1]\n",
    "        global_views = torch.cat(views[:2])\n",
    "        local_views = torch.cat(views[2:])\n",
    "\n",
    "        # Masking\n",
    "        B = len(global_views)\n",
    "        sequence_length = self.teacher_backbone.sequence_length\n",
    "        mask = global_views.new_zeros((B, sequence_length), dtype=torch.bool)\n",
    "        # Mask patches except class token.\n",
    "        H, W = self.teacher_backbone.vit.patch_embed.grid_size\n",
    "        assert (\n",
    "            H * W == sequence_length - 1\n",
    "        ), f\"Unexpected grid size: {H}x{W}, sequence_length {sequence_length}\"\n",
    "        block_mask = random_block_mask(size=(B, H, W), device=mask.device)\n",
    "        mask[:, 1:] = block_mask.flatten(start_dim=1)\n",
    "\n",
    "        # Teacher forward\n",
    "        with torch.no_grad():\n",
    "            teacher_cls_token, teacher_features = self.forward_teacher(global_views)\n",
    "            teacher_cls_out = self.teacher_head.dino_head.forward(teacher_cls_token)\n",
    "            teacher_masked_out = self.teacher_head.ibot_head.forward(\n",
    "                teacher_features[mask]\n",
    "            )\n",
    "\n",
    "        # Student forward\n",
    "        student_global_cls_token, student_global_masked_features = self.forward_student(\n",
    "            global_views, mask=mask\n",
    "        )\n",
    "        student_global_cls_out = self.student_head.dino_head.forward(\n",
    "            student_global_cls_token\n",
    "        )\n",
    "        student_global_masked_out = self.student_head.ibot_head.forward(\n",
    "            student_global_masked_features\n",
    "        )\n",
    "\n",
    "        student_local_cls_token, _ = self.forward_student(local_views, mask=None)\n",
    "        student_local_cls_out = self.student_head.dino_head.forward(\n",
    "            student_local_cls_token\n",
    "        )\n",
    "        student_cls_out = torch.cat([student_global_cls_out, student_local_cls_out])\n",
    "\n",
    "        teacher_temp = linear_warmup_schedule(\n",
    "            step=self.trainer.global_step,\n",
    "            warmup_steps=int(\n",
    "                30 / self.trainer.max_epochs * self.trainer.estimated_stepping_batches\n",
    "            ),\n",
    "            start_value=0.04,\n",
    "            end_value=0.07,\n",
    "        )\n",
    "        dino_loss = self.dino_criterion(\n",
    "            teacher_out=teacher_cls_out.chunk(2),\n",
    "            student_out=student_cls_out.chunk(len(views)),\n",
    "            teacher_temp=teacher_temp,\n",
    "        )\n",
    "        ibot_loss = self.ibot_criterion(\n",
    "            teacher_out=teacher_masked_out,\n",
    "            student_out=student_global_masked_out,\n",
    "            mask=block_mask,\n",
    "            teacher_temp=teacher_temp,\n",
    "        )\n",
    "        koleo_loss = 0.1 * sum(\n",
    "            self.koleo_criterion(t) for t in student_global_cls_token.chunk(2)\n",
    "        )\n",
    "        loss = dino_loss + ibot_loss + koleo_loss\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optim = AdamW(self.parameters(), lr=0.001)\n",
    "        return optim\n",
    "\n",
    "    def on_before_optimizer_step(self, optimizer: AdamW, *args) -> None:\n",
    "        # Optionally zero out the learning rate of the last layer.\n",
    "        if self.current_epoch < 1:\n",
    "            for param_group in optimizer.param_groups:\n",
    "                if \"last_layer\" in param_group:\n",
    "                    param_group[\"lr\"] = 0.0\n",
    "\n",
    "        # Apply weight decay schedule\n",
    "        weight_decay = cosine_schedule(\n",
    "            step=self.trainer.global_step,\n",
    "            max_steps=self.trainer.estimated_stepping_batches,\n",
    "            start_value=0.04,\n",
    "            end_value=0.4,\n",
    "        )\n",
    "        for group in optimizer.param_groups:\n",
    "            if group[\"weight_decay\"] != 0.0:\n",
    "                group[\"weight_decay\"] = weight_decay\n",
    "\n",
    "    def on_train_batch_end(self, outputs, batch, batch_idx):\n",
    "        # Momentum update teacher.\n",
    "        momentum = cosine_schedule(\n",
    "            step=self.trainer.global_step,\n",
    "            max_steps=self.trainer.estimated_stepping_batches,\n",
    "            start_value=0.992,\n",
    "            end_value=1.0,\n",
    "        )\n",
    "        update_momentum(self.student_backbone, self.teacher_backbone, m=momentum)\n",
    "        update_momentum(self.student_head, self.teacher_head, m=momentum)\n",
    "\n",
    "        return super().on_train_batch_end(outputs, batch, batch_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = DINOv2()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = DINOTransform(\n",
    "    global_crop_scale=(0.32, 1),\n",
    "    local_crop_scale=(0.05, 0.32),\n",
    "    n_local_views=8,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we ignore object detection annotations by setting target_transform to return 0\n",
    "def target_transform(t):\n",
    "    return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = torchvision.datasets.VOCDetection(\n",
    "    \"datasets/pascal_voc\",\n",
    "    download=True,\n",
    "    transform=transform,\n",
    "    target_transform=target_transform,\n",
    ")\n",
    "# or create a dataset from a folder containing images or videos:\n",
    "# dataset = LightlyDataset(\"path/to/folder\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = torch.utils.data.DataLoader(\n",
    "    dataset,\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    num_workers=8,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train with DDP and use Synchronized Batch Norm for a more accurate batch norm\n",
    "# calculation. Distributed sampling is also enabled with replace_sampler_ddp=True.\n",
    "trainer = pl.Trainer(\n",
    "    max_epochs=50,\n",
    "    devices=\"auto\",\n",
    "    accelerator=\"gpu\",\n",
    "    strategy=\"ddp_find_unused_parameters_true\",\n",
    "    sync_batchnorm=True,\n",
    "    use_distributed_sampler=True,  # or replace_sampler_ddp=True for PyTorch Lightning <2.0\n",
    ")\n",
    "trainer.fit(model=model, train_dataloaders=dataloader)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
